'''
-----------------------------
This file contains 3 classes:
AgglomerativeClustering,
DivisiveHierarchicalKMeans,
KMeansClustering
-----------------------------
'''

import numpy
import linecache
from ete2 import Tree
import random
import math
import os

'''
---------------------------------------------------------------------------------------
This class is used to cluster objects using distance matrix from ProximityMatrix folder
For more information on obtaining proximities see Proximity file
OUTPUT: a tree written to Tree folder
The algorithm used - agglomerative bottom up merging of clusters
For calculating distances between clusters can use one of three types of distances: 
Maximum, Minimum and UPGMA (weighted average)
USAGE:
cl = Clustering.AgglomerativeClustering('features_lewis37900', 'Manhattan')
cl.Clustering('Maximum')
----------------------------------------------------------------------------------------
'''

class AgglomerativeClustering(object):

    def __init__(self, file, distance_type):
        
        print 'Start initialisation'
        
        self.file = file        
        self.distance_type = distance_type
        self.path1 = os.path.dirname(__file__) + '/ProximityMatrix/'
        self.path2 = os.path.dirname(__file__) + '/Tree/'
        #self.path1 = 'ProximityMatrix/'
        #self.path2 = 'Tree/'
        self.input = self.path1 + file + '_' + distance_type + '.dist'
        def GetNumberofLines():
            lines = 0
            for line in open(self.input):
                lines += 1
            return lines
        self.lines = GetNumberofLines()
        
        self.hierarchical_clustering = Tree()
        self.ProximityMatrix = numpy.zeros((self.lines, self.lines))
        # Index list keeps track of indeces of merged images
        self.Forest = []
        for i in range(self.lines):
            t = Tree()
            t.add_child(name=i)
            self.Forest.append(t)
        
        self.LoadProximityMatrix()
        
    def LoadProximityMatrix(self):
        self.ProximityMatrix = numpy.genfromtxt(self.input, unpack=True)
        for i in range(self.lines):
#            print 'Reading distances = ', i
#            self.ProximityMatrix[i :] = self.GetVector(i+1)
            self.ProximityMatrix[i,i] = numpy.inf
            
    def GetVector(self, number):
        # Get vector from line #number from file input
        # number starts from 1
        line = linecache.getline(self.input, number)
        return map(float, line.split())
    
    def Clustering(self, clustering_type):
        print 'Start clustering'
        i = 1
        while self.ProximityMatrix.size != 1:
            print 'Clustering = ', i
            #print self.ProximityMatrix
            clusters_to_merge, mindistance = self.Merge(clustering_type)
            #print clusters_to_merge
            self.ModifyTree(clusters_to_merge, mindistance)
            i = i+1
        for tree in self.Forest:
            print tree
            tree.write(outfile = self.path2 + self.file + '_' + self.distance_type + '_' + clustering_type + '.nw')
            
    def Merge(self, clustering_type):
        minval = numpy.argmin(self.ProximityMatrix)
        mindistance = numpy.min(self.ProximityMatrix)
        clusters_to_merge = numpy.array(numpy.unravel_index(minval, self.ProximityMatrix.shape))
#        print clusters_to_merge
#        print self.ProximityMatrix
#        print self.CalculateNewDistances(type, clusters_to_merge)
        new_distances = self.CalculateNewDistances(clustering_type, clusters_to_merge)
        self.ProximityMatrix = numpy.vstack((self.ProximityMatrix,new_distances))
        new_distances = numpy.append(new_distances,numpy.inf)
        self.ProximityMatrix = numpy.column_stack((self.ProximityMatrix, new_distances))
#        print self.ProximityMatrix
        self.ProximityMatrix = numpy.delete(self.ProximityMatrix, clusters_to_merge, 0)
        self.ProximityMatrix = numpy.delete(self.ProximityMatrix, clusters_to_merge, 1)
        
        return clusters_to_merge, mindistance
        
    def ModifyTree(self, clusters_to_merge, mindistance):
        if clusters_to_merge[0]>clusters_to_merge[1]:
            cluster1 = self.Forest.pop(clusters_to_merge[0])
            cluster2 = self.Forest.pop(clusters_to_merge[1])
        else:
            cluster1 = self.Forest.pop(clusters_to_merge[1])
            cluster2 = self.Forest.pop(clusters_to_merge[0])
        new_tree = Tree()
        new_tree.add_child(cluster1, dist=mindistance)
        new_tree.add_child(cluster2, dist=mindistance)
        self.Forest.append(new_tree)
        #for i in self.Forest:
         #   print i
    
    def CalculateNewDistances(self, type, clusters_to_merge):
        if type == 'Maximum':
            distance = numpy.max(self.ProximityMatrix[clusters_to_merge], axis=0)
        elif type == 'Minimum':
            distance = numpy.min(self.ProximityMatrix[clusters_to_merge], axis=0)
        elif type == 'UPGMA':
            distance = numpy.mean(self.ProximityMatrix[clusters_to_merge], axis=0)
        return distance    
'''
---------------------------------------------------------------------------------------------------------------
Another class for clustering, it's advantage that it allows to have non-binary structure of classes
Uses KMeansClustering as sub rutine, Silhouette measure is calculated to decide on number of clusters at a level
For calculating Silhouette needs to have a file with pairwise distances
CLUSTERS_NUMBER determines possible size of the cluster
REPEAT determines how many times each clustering will be performed (to avoid wrong clustering due to initialization)
MAXIMAGES how many images at max could be in the cluster so that we try clustering according to CLUSTERS_NUMBER,
otherwise do binary clustering
USAGE:
clustering = DivisiveHierarchicalKMeans(filename)
images = clustering.GetInitialImages()
clustering.TreeCluster(images, range(len(images)))
tree = clustering.hierarchical_clustering
# Wrtie the result to a file
tree.write(outfile = '/home/konyushk/workspace/MCTS/Tree/' + file + '_Euclidean' + '_kmeans.nw') 
----------------------------------------------------------------------------------------------------------------
'''

class DivisiveHierarchicalKMeans(object):
    CLUSTERS_NUMBER = [2, 3, 4, 5]
    REPEAT = 3
    PROGRESS = 0
    MAXIMAGES = 1001
    
    def __init__(self, file):
        self.file = file
        self.input = file
        def GetNumberofLines():
            lines = 0
            for line in open(self.input):
                lines += 1
            return lines
        def GetNumberofColumns():
            line = linecache.getline(self.input, 1)
            return len(line.split())
        self.lines = GetNumberofLines()
        self.columns = GetNumberofColumns()
        self.hierarchical_clustering = Tree()
        for i in range(self.lines):
            self.hierarchical_clustering.add_child(name=i)
           
    def GetInitialImages(self):
        print 'Start reading images'
        images  = numpy.zeros((self.lines, self.columns))
        for i in range(self.lines):
            images[i :] = self.GetVector(i+1)
        print 'Images read'
        return images
            
    def GetVector(self, number):
        line = linecache.getline(self.input, number)
        return map(float, line.split())
        
    def TreeCluster(self, images, image_names):
         
        if len(images) < DivisiveHierarchicalKMeans.MAXIMAGES:   
            silhouette_cn = dict()
            for cn in DivisiveHierarchicalKMeans.CLUSTERS_NUMBER:
                silhouette = numpy.zeros((DivisiveHierarchicalKMeans.REPEAT))
                for i in range(DivisiveHierarchicalKMeans.REPEAT):
                    kmclustering = KMeansClustering(images, cn, self.file)
                    clusters , centroids, silhouette[i] = kmclustering.KMeans()
                silhouette_cn[cn] = numpy.mean(silhouette)
            best_clusters_number = max(silhouette_cn, key=silhouette_cn.get)
            kmclustering = KMeansClustering(images, best_clusters_number, self.file)
            clusters , centroids, silhouette[i] = kmclustering.KMeans()
        else:
            kmclustering = KMeansClustering(images, 2, self.file)
            clusters , centroids = kmclustering.BasicKMeans()
            
        for cluster in clusters:
            if len(clusters[cluster]) >= DivisiveHierarchicalKMeans.CLUSTERS_NUMBER[-1]:
             #   print cluster
                branch = Tree()
                #branch = subtree.add_child()
                NodeAdded = False
                i = 0
                for image in clusters[cluster]:
                    i += 1
               #     print i
                    actual_image_name = image_names[image]              
                    old_node = self.hierarchical_clustering.search_nodes(name = str(actual_image_name))[0]
                    if NodeAdded != True:
                        old_node.add_sister(branch)
                        NodeAdded = True
                    branch.add_child(name = actual_image_name)
                    old_node.delete() 
                new_images = images[clusters[cluster],:]
                image_names = numpy.array(image_names)
                new_image_names = image_names[clusters[cluster]]
                self.TreeCluster(new_images, new_image_names)
                #print len(self.hierarchical_clustering.get_leaves())
            else:
                DivisiveHierarchicalKMeans.PROGRESS += len(clusters[cluster])
                #print DivisiveHierarchicalKMeans.PROGRESS 
                branch = Tree()
                NodeAdded = False
                for image in clusters[cluster]:
                    actual_image_name = image_names[image]
                    '''new'''
                    old_node = self.hierarchical_clustering.search_nodes(name = str(actual_image_name))[0]
                    if NodeAdded != True:
                        old_node.add_sister(branch)
                        NodeAdded = True
                    branch.add_child(name = actual_image_name)
                    old_node.delete() 
                    '''
                    branch.add_child(name = actual_image_name)
                    old_node = self.hierarchical_clustering.search_nodes(name = str(actual_image_name))[0]
                    old_parent = old_node.up
                    old_node.delete() 
                old_parent.add_child(branch)'''
    
            
class KMeansClustering(object):
    
    E = 0.0000001
    
    def __init__(self, images, clusters_number, file):
        
        self.file = file
        self.images = images
        self.clusters_number = clusters_number
        self.columns = numpy.size(images, axis=1)
        self.lines = numpy.size(images, axis=0)
        self.silhouette = 0
        self.InitializeCentroids()
        self.clusters = dict()
        self.clusters_list = []
        
    def InitializeCentroids(self):
            self.centroids = numpy.zeros((self.clusters_number, self.columns))
            self.old_centroids = self.centroids
            centroid_indeces = []
            while len(centroid_indeces) < self.clusters_number:
                centroid_index = random.randint(0, self.lines-1)
                if centroid_index not in centroid_indeces:
                    centroid_indeces.append(centroid_index)
            i = 0
            for centroid_index in centroid_indeces:
                #self.centroids[i,:] = numpy.array(self.GetVector(centroid_index))
                self.centroids[i,:] = self.images[centroid_index, :]
                i += 1 
                
    def BasicKMeans(self): 
#        print 'Basic K-means started!'
        old_SSE = 0
        difference = 1
        while  difference > KMeansClustering.E:
#            print 'Iteration started'
            SSE = self.AssignImages()
#            print 'Assigning images done'
            self.UpdateCentroids()
#            print 'Centroids Updated'
            difference = abs(old_SSE - SSE)
            old_SSE = SSE
#            print 'Centroids = ', self.centroids
#            print 'Clusters = ', self.clusters
#            print 'SSE = ', SSE
#        for cluster in self.clusters:
#            print cluster, ' ', len(self.clusters[cluster])
#        print 'K-means finished!'
            
        return self.clusters, self.centroids
    
    def KMeans(self): 
#        print 'K-means started!'
        old_SSE = 0
        difference = 1
        while  difference > KMeansClustering.E:
#            print 'Iteration started'
            SSE = self.AssignImages()
#            print 'Assigning images done'
            self.UpdateCentroids()
#            print 'Centroids Updated'
            difference = abs(old_SSE - SSE)
            old_SSE = SSE
#            print 'Centroids = ', self.centroids
#            print 'Clusters = ', self.clusters
#            print 'SSE = ', SSE
#        for cluster in self.clusters:
#            print cluster, ' ', len(self.clusters[cluster])
#        print 'K-means finished!'
#        print 'Calculating Silhoette started'
        self.CalculteSilhouette(self.file)
#        print 'Silhouette = ', self.silhouette
            
        return self.clusters, self.centroids, self.silhouette
    
    def AssignImages(self):
        self.clusters.clear()
        image_index = 0
        SSE = 0
        for image in self.images:
            distances = numpy.zeros((1, len(self.centroids)))
            i = 0
            for centroid in self.centroids:
                distances[0, i] = self.Distance('Euclidean', image, centroid)
                i += 1
            self.clusters_list.append(numpy.argmin(distances))
            SSE += numpy.min(distances)**2
            if numpy.argmin(distances) in self.clusters:
                self.clusters[numpy.argmin(distances)].append(image_index) 
            else:
                self.clusters[numpy.argmin(distances)] = [image_index] 
            image_index += 1
            # if one of the clusters ended up with an empty one
        while len(self.clusters) < self.clusters_number:
            new_centroid = random.randint(0, self.lines-1)
            self.clusters[len(self.clusters)] = [new_centroid]
            for cluster in self.clusters:
                if new_centroid in self.clusters[cluster]:
                    self.clusters[cluster].remove(new_centroid)
                    self.clusters[len(self.clusters)] = [new_centroid]
                    break
        return SSE
    
    def Distance(self, type, vector1, vector2):
            if type == 'Euclidean':
                distance = math.sqrt(sum((vector1 - vector2)**2))
            elif type == 'Manhattan':
                distance = sum(abs(vector1-vector2))
            elif type == 'Maximum':
                distance = max(vector1-vector2) 
            elif type == 'Mahalanoibis':
                matrix = numpy.vstack((vector1,vector2)).T
                covariance = numpy.cov(matrix)
                distance = math.sqrt(numpy.dot(numpy.dot((vector1-vector2),numpy.linalg.pinv(covariance)),(vector1-vector2).T))
            elif type == 'Cosine':
                distance = 1-numpy.dot(vector1, vector2)/(math.sqrt(sum((vector1)**2))*math.sqrt(sum((vector2)**2)))
            return distance
        
    def UpdateCentroids(self):
        self.centroids = numpy.zeros((self.clusters_number, self.columns))
        for cluster in self.clusters:
            images_in_clusters = self.images[self.clusters[cluster]]
            self.centroids[cluster, :] = numpy.mean(images_in_clusters, axis=0)
            
    def CalculteSilhouette(self, filename):
        file = filename + '_Euclidean.dist'
        silhouette = []
        for image in range(self.lines):
            distances = numpy.array(self.GetVector(file, image + 1))
            image_cluster = self.clusters_list[image]
            images_in_same_cluster = self.clusters[image_cluster]
            #print image
            #print self.clusters[image_cluster]
            a = numpy.mean(distances[images_in_same_cluster])
            distances_other_clusters = []
            for cluster in self.clusters:
                if cluster != image_cluster:
                    images_in_cluster = self.clusters[cluster]
                    distances_other_clusters.append(numpy.mean(distances[images_in_cluster]))
            b = min(distances_other_clusters)
            silhouette.append((b - a) / max(a, b))
        self.silhouette = sum(silhouette)/len(silhouette)
        
    def GetVector(self, file, number):
        line = linecache.getline(file, number)
        return map(float, line.split())

'''clustering = DivisiveHierarchicalKMeans('test')
images = clustering.GetInitialImages()
clustering.TreeCluster(images, range(len(images)))
tree = clustering.hierarchical_clustering
print tree'''
    
inf = os.path.dirname(__file__) + '-files/test'
clustering = DivisiveHierarchicalKMeans(inf)
images = clustering.GetInitialImages()
clustering.TreeCluster(images, range(len(images)))
tree = clustering.hierarchical_clustering
outf = inf + '_tree.nw'
tree.write(outfile = outf) 